from torch import nn

class NeuralNetwork(nn.Module):

    def __init__(self):
        super().__init__()
        self.lstm = nn.LSTM(input_size=4, hidden_size=4, num_layers=3, batch_first=True)
        self.linear1 = nn.Linear(in_features=4, out_features=3)
        self.tanh1 = nn.Tanh()
        self.linear2 = nn.Linear(in_features=3, out_features=2)
        self.tanh2 = nn.Tanh()
        self.linear_out = nn.Linear(in_features=2, out_features=1)

    def forward(self, input_data):
        lstm_out, _ = self.lstm(input_data)
        l1 = self.tanh1(self.linear1(lstm_out[:, -1, :]))
        l2 = self.tanh2(self.linear2(l1))
        return self.linear_out(l2)

class NeuralNetwork2(nn.Module):

    def __init__(self):
        super().__init__()
        self.lstm = nn.LSTM(input_size=4, hidden_size=4, num_layers=1, batch_first=True)
        self.linear_out = nn.Linear(in_features=4, out_features=3)

    def forward(self, input_data):
        lstm_out, _ = self.lstm(input_data)
        return self.linear_out(lstm_out[:, -1, :])

if __name__ == '__main__':
    from dataset import iter_data, iter_data2
    import torch
    from torch.nn import Softmax

    sm = Softmax(dim=1)

    with torch.no_grad():
        model = NeuralNetwork()
        print(model)
        model.eval()
        for x, y in iter_data(2):
            print('model input:', x)
            print('model output:', model(x))
            print('y:', y)
            break

        model = NeuralNetwork2()
        print(model)
        model.eval()
        for x, y in iter_data2(2):
            print('model input:', x)
            out = model(x)
            print('model output:', out)
            sm_out = sm(out)
            print('class:', sm_out.argmax(1))
            print('y:', y)
            break
